package agent;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Scanner;
import java.util.UUID;
import com.alibaba.fastjson.JSON;
import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.LlmAgent;
import com.google.adk.artifacts.InMemoryArtifactService;
import com.google.adk.events.Event;
import com.google.adk.runner.Runner;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.InMemorySessionService;
import com.google.adk.sessions.Session;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Content;
import com.google.genai.types.Part;
import common.MessageMissionInfo;
import common.QWModel;
import common.QWModelRegistry;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import org.apache.rocketmq.client.apis.ClientConfiguration;
import org.apache.rocketmq.client.apis.ClientException;
import org.apache.rocketmq.client.apis.ClientServiceProvider;
import org.apache.rocketmq.client.apis.SessionCredentialsProvider;
import org.apache.rocketmq.client.apis.StaticSessionCredentialsProvider;
import org.apache.rocketmq.client.apis.consumer.ConsumeResult;
import org.apache.rocketmq.client.apis.consumer.LitePushConsumer;
import org.apache.rocketmq.client.apis.message.Message;
import org.apache.rocketmq.client.apis.producer.Producer;
import org.apache.rocketmq.client.apis.producer.ProducerBuilder;
import org.apache.rocketmq.client.apis.producer.SendReceipt;
import org.apache.rocketmq.client.apis.producer.TransactionChecker;
import org.apache.rocketmq.shaded.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SupervisorAgentMain {
    private static final Logger log = LoggerFactory.getLogger(SupervisorAgentMain.class);
    private static final String GlobalEndpoint = System.getProperty("endpoint");
    private static final String RocketMQInstanceID = System.getProperty("rocketmqInstanceID");
    private static final String WorkAgentResponseTopic = System.getProperty("workAgentResponseTopic");
    private static final String WorkAgentResponseGroupID = System.getProperty("workAgentResponseGroupID");
    private static final String AccessKey = System.getProperty("ak");
    private static final String SecretKey = System.getProperty("sk");
    private static final String ApiKey = System.getProperty("apiKey");
    private static final String WeatherAgentTaskTopic = System.getProperty("weatherAgentTaskTopic");
    private static final String TravelAgentTaskTopic = System.getProperty("travelAgentTaskTopic");
    private static final String AgentName = "SupervisorAgent";

    private static final String USER_ID = "rocketmq_a2a_user";
    private static final String APP_NAME = "rocketmq_a2a";
    private static final Map<String/* topic */, Producer> topicProducerMap = new HashMap<>();
    private static String lastQuestion = "";

    public static void main(String[] args) throws Exception {
        if (!checkConfigParam()) {
            System.out.println("配置参数不完整，请检查参数配置情况");
            return;
        }
        BaseAgent baseAgent = initAgent(WeatherAgentTaskTopic, TravelAgentTaskTopic);
        printSystemInfo("🚀 启动 QWen为底座模型的 " + AgentName + "，擅长处理天气问题与行程安排规划问题，在本例中使用RocketMQ LiteTopic实现多个Agent之间的通讯");
        printSystemInfo("📋 初始化会话...");
        InMemoryArtifactService artifactService = new InMemoryArtifactService();
        InMemorySessionService sessionService = new InMemorySessionService();
        Runner runner = new Runner(baseAgent, APP_NAME, artifactService, sessionService, /* memoryService= */ null);
        Session session = runner
            .sessionService()
            .createSession(APP_NAME, USER_ID)
            .blockingGet();
        LitePushConsumer litePushConsumer = buildConsumer(sessionService, runner);
        subLiteTopic(litePushConsumer, session.id());
        printSystemSuccess("✅ 会话创建成功: " + session.id());
        printSystemInfo("💡 输入 'quit' 退出，输入 'help' 查看帮助");
        try (Scanner scanner = new Scanner(System.in, "UTF-8")) {
            while (true) {
                printPrompt("You");
                String userInput = scanner.nextLine().trim();
                if ("quit".equalsIgnoreCase(userInput)) {
                    printSystemInfo("👋 再见！");
                    System.exit(0);
                    break;
                }
                if ("help".equalsIgnoreCase(userInput)) {
                    printHelp();
                    continue;
                }
                if (StringUtils.isEmpty(userInput)) {
                    printSystemInfo("请不要输入空值");
                    continue;
                }
                printSystemInfo("🤔 正在思考...");
                log.info("用户输入: {}", userInput);
                Content userMsg = Content.fromParts(Part.fromText(userInput));
                Flowable<Event> events = runner.runAsync(USER_ID, session.id(), userMsg);
                events.blockingForEach(event -> {
                    String content = event.stringifyContent();
                    dealEventContent(content, session);
                });
            }
        }
    }

    private static boolean checkConfigParam() {
        if (StringUtils.isEmpty(GlobalEndpoint) || StringUtils.isEmpty(RocketMQInstanceID)
            || StringUtils.isEmpty(WorkAgentResponseTopic) || StringUtils.isEmpty(WorkAgentResponseGroupID)
            || StringUtils.isEmpty(AccessKey) || StringUtils.isEmpty(SecretKey)
            || StringUtils.isEmpty(ApiKey) || StringUtils.isEmpty(WeatherAgentTaskTopic) || StringUtils.isEmpty(
            TravelAgentTaskTopic)) {
            if (StringUtils.isEmpty(GlobalEndpoint)) {
                System.out.println("请配置RocketMQV5的实例接入点 endpoint");
            }
            if (StringUtils.isEmpty(RocketMQInstanceID)) {
                System.out.println("请配置RocketMQV5的实例信息 rocketmqInstanceID");
            }
            if (StringUtils.isEmpty(WorkAgentResponseTopic)) {
                System.out.println("请配置RocketMQV5的轻量消息Topic workAgentResponseTopic");
            }
            if (StringUtils.isEmpty(WorkAgentResponseGroupID)) {
                System.out.println("请配置RocketMQV5的轻量消息消费者 workAgentResponseGroupID");
            }
            if (StringUtils.isEmpty(AccessKey)) {
                System.out.println("请配置RocketMQV5的访问控制-用户名 ak");
            }
            if (StringUtils.isEmpty(SecretKey)) {
                System.out.println("请配置RocketMQV5的访问控制-密码 sk");
            }
            if (StringUtils.isEmpty(ApiKey)) {
                System.out.println("请配置SupervisorAgent qwen-plus 在百炼平台的 apiKey");
            }
            if (StringUtils.isEmpty(WeatherAgentTaskTopic)) {
                System.out.println("请配置RocketMQV5的普通消息Topic weatherAgentTaskTopic");
            }
            if (StringUtils.isEmpty(TravelAgentTaskTopic)) {
                System.out.println("请配置RocketMQV5的普通消息Topic travelAgentTaskTopic");
            }
            return false;
        }
        return true;
    }

    private static void dealEventContent(String content, Session session) {
        if (content != null && !content.trim().isEmpty()) {
            if (content.startsWith("{")) {
                try {
                    Mission mission = JSON.parseObject(content, Mission.class);
                    if (null != mission) {
                        printPrompt("Agent");
                        System.out.println(AgentName + "转发请求到其他的Agent, 等待其他Agent响应，AgentTopic: " + mission.getTopic() + " 问题: " + mission.getMessageInfo());
                        dealMissionByMessage(mission, session.id());
                    }
                } catch (Exception e) {
                    System.out.println("解析过程出现异常");
                }
            } else {
                printPrompt("Agent");
                System.out.println(content);
            }
        }
    }

    private static LitePushConsumer buildConsumer(BaseSessionService baseSessionService, Runner runner)
        throws ClientException {
        final ClientServiceProvider provider = ClientServiceProvider.loadService();
        SessionCredentialsProvider sessionCredentialsProvider = new StaticSessionCredentialsProvider(AccessKey, SecretKey);
        ClientConfiguration clientConfiguration = ClientConfiguration.newBuilder()
            .setEndpoints(GlobalEndpoint)
            .setNamespace(RocketMQInstanceID)
            .setCredentialProvider(sessionCredentialsProvider)
            .build();
        LitePushConsumer litePushConsumer = provider.newLitePushConsumerBuilder()
            .setClientConfiguration(clientConfiguration)
            .setConsumerGroup(WorkAgentResponseGroupID)
            .bindTopic(WorkAgentResponseTopic)
            .setMessageListener(messageView -> {
                try {
                    Optional<String> liteTopic = messageView.getLiteTopic();
                    String liteTopicStr = liteTopic.get();
                    if (StringUtils.isEmpty(liteTopicStr)) {
                        return ConsumeResult.SUCCESS;
                    }
                    byte[] result = new byte[messageView.getBody().remaining()];
                    messageView.getBody().get(result);
                    String resultStr = new String(result, StandardCharsets.UTF_8);
                    MessageMissionInfo messageMissionInfo = JSON.parseObject(resultStr, MessageMissionInfo.class);
                    if (null == messageMissionInfo) {
                        return ConsumeResult.SUCCESS;
                    }
                    printPrompt("Agent");
                    System.out.println("收到" + messageMissionInfo.getResponseAgentName() + " 返回的响应结果, 之前提出的问题:" + messageMissionInfo.getMessageRequest());
                    System.out.println("响应结果: " + messageMissionInfo.getMessageText());
                    printSystemSuccessSimple("响应结果来自于RocketMQ消息, WorkAgentResponseTopic: " + WorkAgentResponseTopic + ", LiteTopic: " + liteTopicStr + ", messageId: " + messageView.getMessageId());
                    Maybe<Session> sessionMaybe = baseSessionService.getSession(APP_NAME, USER_ID, liteTopicStr, Optional.empty());
                    Event event = Event.builder()
                        .id(UUID.randomUUID().toString())
                        .invocationId(UUID.randomUUID().toString())
                        .author(APP_NAME)
                        .content(buildContent(messageMissionInfo.getMessageText()))
                        .build();
                    Session session = sessionMaybe.blockingGet();
                    baseSessionService.appendEvent(session, event);
                    Content userMsg = Content.fromParts(Part.fromText(messageMissionInfo.getMessageText()));
                    Flowable<Event> events = runner.runAsync(USER_ID, session.id(), userMsg);
                    events.blockingForEach(eventSub -> {
                        boolean equals = lastQuestion.equals(eventSub.stringifyContent());
                        if (equals) {
                            return;
                        }
                        lastQuestion = eventSub.stringifyContent();
                        String content = lastQuestion;
                        if (content != null && !content.trim().isEmpty()) {
                            if (content.startsWith("{")) {
                                try {
                                    Mission mission = JSON.parseObject(content, Mission.class);
                                    if (null != mission && !StringUtils.isEmpty(mission.getMessageInfo())
                                        && !StringUtils.isEmpty(mission.getTopic())) {
                                        printPrompt("Agent");
                                        System.out.println("转发到其他的Agent, 等待其他Agent响应，AgentTopic: " + mission.getTopic() + " 问题: " + mission.getMessageInfo());
                                        dealMissionByMessage(mission, session.id());
                                    }
                                } catch (Exception e) {
                                    System.out.println("解析过程出现异常");
                                }
                            }
                        } else {
                            System.out.println(content);
                            log.debug("Agent 响应: {}", content);
                        }
                    });
                    printPrompt("You");
                    return ConsumeResult.SUCCESS;
                } catch (Exception e) {
                    System.out.println("拉取消息消费过程中出现异常 messageId: " + messageView.getMessageId() + " error: " + e.getMessage());
                    return ConsumeResult.FAILURE;
                }
            }).build();
        return litePushConsumer;
    }

    private static void subLiteTopic(LitePushConsumer litePushConsumer, String liteTopic) throws ClientException {
        if (null == litePushConsumer || StringUtils.isEmpty(liteTopic)) {
            return;
        }
        litePushConsumer.subscribeLite(liteTopic);
        System.out.println("完成对 LiteTopic: " + liteTopic + "的订阅");
    }

    private static Content buildContent(String content) {
        if (StringUtils.isEmpty(content)) {
            return null;
        }
        return Content.builder()
            .role(APP_NAME)
            .parts(ImmutableList.of(Part.builder().text(content).build()))
            .build();
    }

    private static void dealMissionByMessage(Mission mission, String sessionId) {
        if (null == mission || StringUtils.isEmpty(mission.getTopic()) || StringUtils.isEmpty(mission.getMessageInfo()) || StringUtils.isEmpty(sessionId)) {
            return;
        }
        try {
            final ClientServiceProvider provider = ClientServiceProvider.loadService();
            String topic = mission.getTopic().replaceAll(" ", "");
            Producer producer = null;
            try {
                producer = getOrCreateProducer(topic);
            } catch (ClientException e) {
                return;
            }
            if (null == producer) {
                System.out.println("producer is null, topic: " + topic);
                return;
            }
            String bodyStr = serialText(mission.getMessageInfo(), sessionId);
            if (StringUtils.isEmpty(bodyStr)) {
                System.out.println("bodyStr is null, topic: " + topic);
                return;
            }
            byte[] body = bodyStr.getBytes(StandardCharsets.UTF_8);
            final Message message = provider.newMessageBuilder()
                .setTopic(topic)
                .setBody(body)
                .build();
            try {
                final SendReceipt sendReceipt = producer.send(message);
                printSystemSuccess("发送消息到Agent的普通的Topic 成功, topic: " + topic + ", messageId: " + sendReceipt.getMessageId() + ", sessionId: " + sessionId);
            } catch (Throwable t) {
                System.out.println("发送消息异常");
            }
        } catch (Exception e) {
            System.out.println("error" + e.getMessage());
        }
    }

    private static Producer getOrCreateProducer(String topic) throws ClientException {
        if (StringUtils.isEmpty(topic)) {
            return null;
        }
        Producer producer = topicProducerMap.get(topic);
        if (null != producer) {
            return producer;
        }
        producer = buildProducer(null, topic);
        topicProducerMap.put(topic, producer);
        return producer;
    }

    private static Producer buildProducer(TransactionChecker checker, String... topics) throws ClientException {
        final ClientServiceProvider provider = ClientServiceProvider.loadService();
        SessionCredentialsProvider sessionCredentialsProvider = new StaticSessionCredentialsProvider(AccessKey, SecretKey);
        ClientConfiguration clientConfiguration = ClientConfiguration.newBuilder()
            .setEndpoints(GlobalEndpoint)
            .setNamespace(RocketMQInstanceID)
            .setCredentialProvider(sessionCredentialsProvider)
            .setRequestTimeout(Duration.ofSeconds(10))
            .build();
        final ProducerBuilder builder = provider.newProducerBuilder()
            .setClientConfiguration(clientConfiguration)
            .setTopics(topics);
        if (checker != null) {
            builder.setTransactionChecker(checker);
        }
        return builder.build();
    }

    private static String serialText(String messageText, String liteTopic) {
        if (StringUtils.isEmpty(messageText) || StringUtils.isEmpty(liteTopic)) {
            return null;
        }
        MessageMissionInfo messageMissionInfo = new MessageMissionInfo();
        messageMissionInfo.setMessageText(messageText);
        messageMissionInfo.setLiteTopic(liteTopic);
        messageMissionInfo.setRequestAgentName(AgentName);
        return JSON.toJSONString(messageMissionInfo);
    }

    private static void printSystemInfo(String message) {
        System.out.println("\u001B[34m[SYSTEM] " + message + "\u001B[0m");
        log.info(message);
    }

    private static void printSystemSuccess(String message) {
        System.out.println("\u001B[32m[SUCCESS] " + message + "\u001B[0m");
        log.info(message);
    }

    private static void printSystemSuccessSimple(String message) {
        System.out.print("\u001B[32m[SUCCESS] " + message + "\u001B[0m");
    }

    private static void printPrompt(String role) {
        System.out.print("\n\u001B[36m" + role + " > \u001B[0m");
    }

    private static void printHelp() {
        System.out.println("\n\u001B[35m📖 帮助信息:\u001B[0m");
        System.out.println("  • 询问天气: '杭州下周的天气情况怎么样'");
        System.out.println("  • 帮忙安排行程: '帮我做一个下周三到下周日杭州周边自驾游方案'");
        System.out.println("  • 退出程序: 'quit'");
        System.out.println("  • 显示帮助: 'help'");
    }

    public static BaseAgent initAgent(String weatherTopic, String travelTopic) {
        if (StringUtils.isEmpty(weatherTopic) || StringUtils.isEmpty(travelTopic)) {
            System.out.println("initAgent 参数缺失，请补充天气Topic、旅行安排Topic");
            return null;
        }
        QWModel qwModel = QWModelRegistry.getModel(ApiKey);
        return LlmAgent.builder()
            .name(APP_NAME)
            .model(qwModel)
            .description("你是一位专业的行程规划专家")
            .instruction("# 角色\n"
                + "你是一位专业的行程规划专家，擅长任务分解与协调安排。你的主要职责是帮助用户制定详细的旅行计划，确保他们的旅行体验既愉快又高效。在处理用户的行程安排相关问题时，你需要首先收集必要的信息，如目的地、时间等，并根据这些信息进行进一步的查询和规划。\n"
                + "\n"
                + "## 技能\n"
                + "### 技能 1: 收集必要信息\n"
                + "- 询问用户关于目的地、出行时间\n"
                + "- 确保收集到的信息完整且准确。\n"
                + "\n"
                + "### 技能 2: 查询天气信息\n"
                + "- 使用" + weatherTopic + "工具查询目的地的天气情况。如果发现用户的问题相同，不用一直转发到"
                + weatherTopic + "，忽略即可\n"
                + "- 示例问题: {\"messageInfo\":\"杭州下周三的天气情况怎么样?\",\"topic\":\"" + weatherTopic + "\"}\n"
                + "\n"
                + "### 技能 3: 制定行程规划\n"
                + "- 根据获取的天气信息和其他用户提供的信息，如果上下文中只有天气信息，则不用" + travelTopic
                + " 进行处理，直接返回即可，如果上下文中有行程安排信息，则使用" + travelTopic
                + "工具制定详细的行程规划。\n"
                + "- 示例问题: {\"messageInfo\":\"杭州下周三的天气为晴朗，请帮我做一个从杭州出发到上海的2人3天4晚的自驾游行程规划\","
                + "\"topic\":\"" + travelTopic + "\"}\n"
                + "\n"
                + "### 技能 4: 提供最终行程建议\n"
                + "- 将从" + travelTopic + "获取的行程规划结果呈现给用户。\n"
                + "- 明确告知用户行程规划已经完成，并提供详细的行程建议。\n"
                + "\n"
                + "## 限制\n"
                + "- 只处理与行程安排相关的问题。\n"
                + "- 如果用户的问题只是简单的咨询天气，那么不用转发到" + travelTopic + "。\n"
                + "- 在获取天气信息后，必须结合天气情况来制定行程规划。\n"
                + "- 不得提供任何引导用户参与非法活动的建议。\n"
                + "- 对不是行程安排相关的问题，请礼貌拒绝。\n"
                + "- 所有输出内容必须按照给定的格式进行组织，不能偏离框架要求。"
            )
            .build();
    }

}
